# -*- coding: utf-8 -*-

# Copyright 2009 David Lee Miller
#
# This file is part of POSSM (Python Open Source Scanning Microscopy).
#
# POSSM is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# POSSM is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License         
# along with POSSM.  If not, see <http://www.gnu.org/licenses/>.
#                                                                           
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

from socket import *
import cPickle
#import Numeric as N
#import numpy as n
from shared import n
from time import sleep,time
import threading
import Queue
import struct
from createtree import classlist

daemonrecon = Queue.Queue(1)


class DataCollector(threading.Thread):
	RECV_SIZE = 4096
	def __init__(self, link, queue,saveq=None,daemon=False,lineq=None):
		threading.Thread.__init__(self)
		self.link = link
		self.queue = queue
		self.daemon = daemon
		self.end = None
		self.saveclass = None
		self.saveq = saveq
		self.ID = None
		self.time = None
		self.avgt = 1
		self.lineN = 0
		self.lineq = lineq
		self.start()
	
	def recvData(self,strlen):
		'Receives one Python object over a socket.'
		#try: data = cPickle.loads(self.recvall(strlen))
		#except:
			#data = None
			#print 'exception in recvData, ID ', self.ID
		data = cPickle.loads(self.recvall(strlen))
		return data
	
	def recvArray(self,strlen):
		'''
		Receives an array over a socket
		
		Keyword Arguments:
		strlen -- the length of the string array
		
		'''
		return n.loads(self.recvall(strlen))
	
	def recvall(self,strlen):
		'Receives an entire string over a socket.'
		recvd = 0
		string = str()
		while recvd < strlen:
			if strlen-recvd < self.RECV_SIZE: buf_sz = strlen-recvd
			else: buf_sz = self.RECV_SIZE
			newstr = self.link.SCK.recv(buf_sz)
			string = string + newstr
			recvd = recvd + len(newstr)
		return string
	
	def countLine(self):
		if not self.daemon and self.lineq:
			self.lineN += 1
			if self.lineq.empty():
				self.lineq.put(self.lineN)
			else:
				self.lineq.get()
				self.lineq.put(self.lineN)
		else: pass
	
	def avgTime(self):
		if not self.daemon and self.lineq:
			endtime = time()
			self.avgt = (endtime - self.time)/self.lineN
		else: pass
		
	def reconnect(self):
		sleep(0.2)
		try:
			self.link.SCK.close()
			self.link.server()
		except: self.link.server()
	
	def getLen(self):
		if not self.daemon:
			try: 
				strlen = struct.unpack('I',self.link.SCK.recv(4))[0]
				com = self.link.SCK.recv(1)
			except:
				strlen = None
				com = None
		else:
			try:
				strlen = struct.unpack('I',self.link.SCK.recv(4))[0]
				com = self.link.SCK.recv(1)
			except:
				self.reconnect()
				com = None
				strlen = None
		return com,strlen
	
	def Collect(self):
		while 1:
			if self.end: break
			com,strlen = self.getLen()
			if self.end: break
			
			if com == 'a':
				data = self.recvArray(strlen)
				if not self.time: self.time = time()
				self.countLine()
				self.avgTime()
				self.queue.put(data)
			if com == 'b':
				data = self.recvData(strlen)
				if not self.time: self.time = time()
				self.countLine()
				self.avgTime()
				self.queue.put(data)
			elif com == 'c':
				data = self.recvData(strlen)
				#print type(data)
				if type(data) in classlist:
					self.saveclass = data
					if self.saveq: self.saveq.put(data)
				#elif type(data) == type('0'): print data
				self.queue.put(data)
			elif com == 'm':
				#print 'RECEIVING MESSAGE FROM THE GIANT HEAD'
				data = self.recvData(strlen)
				self.queue.put(data)
			elif com == 'p':
				if self.end: break
				else: pass
			elif com == '' and self.daemon:
				try:
					if self.end: break
					self.link.SCK.close()
					self.link.server()
				except: self.link.server()
			elif not self.link: break
			com = None
	
	def run(self):
		if self.daemon: self.link.server()
		self.Collect()
		print 'quitting collector'

class DataDispenser(threading.Thread):
	SEND_SIZE = 1024
	def __init__(self, link, queue, daemon=False):
		threading.Thread.__init__(self)
		self.link = link
		self.queue = queue
		self.daemon = daemon
		self.ID = None
		self.end = None
		self.start()
		
	def sendData(self, data, flag):
		'''Sends one Python object over a socket.'''
		try:
			sdata = cPickle.dumps(data)
			strlen = struct.pack('I',len(sdata))
			self.link.SCK.send(strlen)
			self.link.SCK.send(flag)
			self.link.SCK.sendall(sdata)
		except: print 'exception raised while trying to send class in ', self.ID
	
	def sendArray(self,data):
		try: 
			#data = n.uint32(data)
			sdata = data.dumps()
			strlen = struct.pack('I',len(sdata))
			self.link.SCK.send(strlen)
			self.link.SCK.send('a')
			self.link.SCK.send(sdata)
		except: print 'exception raised while trying to send array in ', self.ID
	
	def sendPass(self):
		try: 
			self.link.SCK.send(struct.pack('I',1))
			self.link.SCK.send('p')
		except:
			if self.daemon: self.link.server()
			else: pass
	
	def Dispense(self):
		typear = type(n.zeros((1)))
		while 1:
			if self.end: break
			try: data = self.queue.get(timeout=1)
			except: data = 'pass'
			
			if data == 'pass': 
				self.sendPass()
				if self.end: break
			elif type(data) == str:
				self.sendData(data,'m')
			elif type(data) == typear: 
				self.sendArray(data)
			elif type(data) in classlist: 
				self.sendData(data,'c')
			elif 'LINE' in dir(data):
				self.sendData(data,'b')
			elif not self.link: break
			else: pass
	def run(self):
		if self.daemon: self.link.server()
		self.Dispense()
		print 'quitting dispenser'

class CmdCollector(DataCollector):
	def __init__(self, link, queue, dae):
		self.cmdq = queue
		DataCollector.__init__(self, link, None, daemon=dae, lineq=None)
		#self.cclass = console_class
	
	def postMess(self,mess):
		self.cmdq.put(mess)
	
	def Collect(self):
		while 1:
			if self.end: break
			com,strlen = self.getLen()
			if com == 'm':
				mess = self.recvData(strlen)
				self.postMess(mess)
			elif com == 'r':
				macro = self.recvData(strlen)
				macro._run(self.cclass)
			elif com == 's':
				macro = self.recvData(strlen)
				self.cmdq.put(macro)
			elif com == 'p':
				if self.end: break
				else: pass
			elif com == '' and self.daemon:
				try:
					if self.end: break
					self.link.SCK.close()
					self.link.server()
				except: self.link.server()
			elif not self.link: break
			com = None

class CmdDispenser(DataDispenser):
	def __init__(self, link, queue, dae):
		DataDispenser.__init__(self, link, queue, daemon=dae)
	
	def Dispense(self):
		while 1:
			if self.end: break
			try: data = self.queue.get(timeout=1)
			except: data = 'pass'
			if data == 'pass': 
				self.sendPass()
			elif type(data) == type('0'):
				try: self.sendData(data,'m')
				except: pass
			elif '_run' in dir(data): 
				self.sendData(data,'r')
			elif 'switchlist' in dir(data): 
				self.sendData(data,'s')

def CloseHandler(link,linkclass):
	try: 
		try:
			#link.connect(gethostname())
			linkclass.end = True
			link.disconnect()
		except: linkclass.end = True
		link = None
	except: print 'Collector Not Closed'


#def CloseDispenser(link, linkclass):
	#try:
		#try:
			#link.connect(gethostname())
			#linkclass.end = True
			#link.disconnect()
		#except: linkclass.end = True
		#link = None
	#except: print 'Dispenser Not Closed'

class ForwardData(threading.Thread):
	'Pass data from one open socket to another'
	
	def __init__(self, link1, link2):
		self.link1 = link1
		self.link2 = link2
		self.queue = Queue.Queue()
		
	def Forwarder(self):
		'''
		Combines a data collector with a data dispenser and uses
		a shared queue to collect data on one socket and pass it directly
		to another socket.
		
		Keyword Arguemnts:
		None
		
		'''
		getter = DataCollector(self.link1,self.queue)
		sender = DataDispenser(self.link2,self.queue)


#This file creates a class for using the socket
#module. This is only intended as a set of basic
#commands and will be subclassed for added functionality
#in specific applications. DLM

class SocketCommands(object):
	'Basic commands for a socket instance'
	SEND_SIZE = 1024
    	RECV_SIZE = 4096
	def __init__(self, port=22222, buf=4096):
		self.PORT = port
		self.BUF = buf
		#self.SCK = socket(AF_INET, SOCK_STREAM)
	
	def connect(self, hostname):
		'''
		Connects to a given host as a client.
		
		Keyword Arguments:
		hostname -- the name (not ip) of the computer to which you wish to connect
		
		'''
		self.SCK = socket(AF_INET, SOCK_STREAM)
		self.HOST = gethostbyname(hostname)
		self.SCK.connect((self.HOST,self.PORT))
	
	def disconnect(self):
		'''
		Disconnects client from host socket.
		
		Keyword Arguments:
		None
		
		'''
		#try: self.SCK.send('end',timeout=0.1)
		#except: 
			#try: self.SCK.sendData('end')
			#except: pass
		self.SCK.close()
	
	def server(self):
		'''
		Creates a server socket and waits for a client connection.
		
		Keyword Arguments:
		None
		
		'''
		self.HOST = gethostname()
		self.SSCK = socket(AF_INET, SOCK_STREAM)
		self.SSCK.bind((self.HOST,self.PORT))
		self.SSCK.listen(1)
		print 'waiting to connect..'
		while 1:
			(self.csock, self.caddr) = self.SSCK.accept()
			self.SCK = self.csock
			if self.caddr: break
		print 'connection received from: ', self.caddr
	
	def sendCom(self, com):
		'''
		Sends a command to the listening party
		
		Keyword Arguments:
		com -- the string command to be sent
		
		'''
		self.SCK.send(com)
		sleep(0.02)
	
	def recvCom(self):
		'''
		Listens for a command from the sending party.
		
		Keyword Arguments:
		None
		
		'''
		self.command = self.SCK.recv(self.BUF)
		return self.command
	
	def sendData(self, dat):
		'''
		Sends data of any object type.
		
		Keyword Arguments:
		dat -- data to be sent
		
		'''
		self.SCK.setblocking(True)
		file = self.SCK.makefile('a',0)
		cPickle.dump(dat, file)
		file.close()
	
	def recvData(self):
		'''
		Listens for data from the sending party.
		
		Keyword Arguments:
		None
		
		'''
		file = self.SCK.makefile()
		self.data = cPickle.load(file)
		print 'received data'
		file.close()
		return self.data
	
	def closeClient(self):
		self.SCK.close()
	
	def closeServer(self):
		self.SSCK.close()
		